package com.example.demo.gateway.filter;

import com.example.demo.gateway.entity.GatewayUpstream;
import lombok.extern.slf4j.Slf4j;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.cloud.gateway.filter.ReactiveLoadBalancerClientFilter;
import org.springframework.cloud.gateway.filter.RouteToRequestUrlFilter;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
import org.springframework.core.Ordered;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.util.UriComponentsBuilder;
import reactor.core.publisher.Mono;

import java.net.URI;
import java.util.Objects;
/**
 * 自定义全局处理器
 * @author zhukaiyuan
 * @Date 2020-05-23 15:26
 */
@Slf4j
@Component
public class  UpstreamLoadBalancerClientGatewayFilter implements GlobalFilter, Ordered {
    private static final String SCHEME_PREFIX = "upstream";
    public static final int UPSTREAM_LOAD_BALANCER_CLIENT_FILTER_ORDER = RouteToRequestUrlFilter.ROUTE_TO_URL_FILTER_ORDER + 1;

    private final UpstreamReactorLoadBalancer loadBalancer;

    public UpstreamLoadBalancerClientGatewayFilter(UpstreamReactorLoadBalancer loadBalancer) {
        this.loadBalancer = loadBalancer;
    }

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        URI url = (URI) exchange.getAttribute(ServerWebExchangeUtils.GATEWAY_REQUEST_URL_ATTR);
        String schemePrefix = (String) exchange.getAttribute(ServerWebExchangeUtils.GATEWAY_SCHEME_PREFIX_ATTR);
        if (url != null && (SCHEME_PREFIX.equals(url.getScheme()) || SCHEME_PREFIX.equals(schemePrefix))) {
            //如果是以upstream则走我们的自定义的负载均衡策略
            if (log.isTraceEnabled()) {
                log.trace(ReactiveLoadBalancerClientFilter.class.getSimpleName() + " url before: " + url);
            }
            String serviceId = url.getHost();
            GatewayUpstream choose = loadBalancer.choose(serviceId);
            URI uri = this.reconstructURI(choose, url);
            exchange.getAttributes().put(ServerWebExchangeUtils.GATEWAY_REQUEST_URL_ATTR, uri);
            return chain.filter(exchange);
        } else {
            return chain.filter(exchange);
        }
    }


    @Override
    public int getOrder() {
        return UPSTREAM_LOAD_BALANCER_CLIENT_FILTER_ORDER;
    }


    protected URI reconstructURI(GatewayUpstream u, URI original) {
        String host = u.getHost();
        String scheme = u.getProtocol();
        int port = u.getPort();
        if (Objects.equals(host, original.getHost()) && port == original.getPort() && Objects.equals(scheme, original.getScheme())) {
            return original;
        } else {
            boolean encoded = containsEncodedParts(original);
            return UriComponentsBuilder.fromUri(original).scheme(scheme).host(host).port(port).build(encoded).toUri();
        }
    }



    private static boolean containsEncodedParts(URI uri) {
        boolean encoded = uri.getRawQuery() != null && uri.getRawQuery().contains("%") || uri.getRawPath() != null && uri.getRawPath().contains("%") || uri.getRawFragment() != null && uri.getRawFragment().contains("%");
        if (encoded) {
            try {
                UriComponentsBuilder.fromUri(uri).build(true);
                return true;
            } catch (IllegalArgumentException var3) {
                return false;
            }
        } else {
            return false;
        }
    }
}
